import librosa
import matplotlib.pyplot as plt
import numpy as np

wave_path=r"wave_3.wav"
waveform, sample_rate=librosa.load(wave_path, sr=None)

def Calc_ZCR(waveform, frame_length, hop_length):
    if len(waveform) % hop_length !=0:
        frame_num = int((len(waveform)-frame_length)/hop_length)+1;
        pad_num=frame_num*hop_length+frame_length-len(waveform)
        #waveform=np.pad(waveform, pad_width=(0, pad_num), mode="wrap")
        waveform=np.pad(waveform, (0, pad_num), 'constant', constant_values=(0, 0))
    print("after padding:", len(waveform))
    frame_num = int((len(waveform)-frame_length)/hop_length)+1
    print("frame_num:", frame_num)
    waveform_zcr=[]
    for t in range(frame_num):#0~61
        #print(t)
        current_frame = waveform[(frame_length-hop_length)*t:(frame_length-hop_length)*t+frame_length]
        #if t==61: #the last frame
            #print(current_frame*32767)
        a=np.sign(current_frame[0: frame_length-1, ])
        b=np.sign(current_frame[1: frame_length, ])
        current_zcr=np.sum(np.abs(a-b))/2/frame_length
        waveform_zcr.append(current_zcr)
    return np.array(waveform_zcr)

def Calc_STE(waveform, frame_length, hop_length):
    if len(waveform) % hop_length !=0:
        frame_num = int((len(waveform)-frame_length)/hop_length)+1;
        pad_num=frame_num*hop_length+frame_length-len(waveform)
        #waveform=np.pad(waveform, pad_width=(0, pad_num), mode="wrap")
        waveform=np.pad(waveform, (0, pad_num), 'constant', constant_values=(0, 0))
    frame_num = int((len(waveform)-frame_length)/hop_length)+1
    print("frame_num:", frame_num)
    waveform_ste=[]
    for t in range(frame_num):
        current_frame = waveform[(frame_length-hop_length)*t:(frame_length-hop_length)*t+frame_length]
        current_ste= sum([element**2 for element in current_frame])
        waveform_ste.append(current_ste)
    return np.array(waveform_ste)

frame_size=512
hop_size=int(frame_size/2)
print("len:", len(waveform))
print("all frame:", int((len(waveform)-frame_size)/hop_size)+1)
print("type:", type(waveform))

waveform_ZCR=Calc_ZCR(waveform, frame_size, hop_size)
#print("done")
#exit() 

print("len-waveform_ZCR",len(waveform_ZCR))

frame_scale=np.arange(0, len(waveform_ZCR), step=1)
print(frame_scale)#0 1 2 3 4 5 ... 61
time_scale=librosa.frames_to_time(frame_scale, hop_length=hop_size)
plt.figure(figsize=(20,10))
plt.plot(time_scale, waveform_ZCR, color="r")
plt.title("Zero-Cross-rate")
#plt.show()

waveform_STE=Calc_STE(waveform, frame_size, hop_size)
frame_scale=np.arange(0, len(waveform_STE), step=1)
time_scale=librosa.frames_to_time(frame_scale, hop_length=hop_size)
plt.figure(figsize=(20,10))
plt.plot(time_scale, waveform_STE, color="r")
plt.title("Zero-short-time-energy")
#plt.show()

#librosa.feature
waveform_zcr_librosa = librosa.feature.zero_crossing_rate(y=waveform, frame_length=frame_size, hop_length=hop_size).T[1:,0]
plt.figure(figsize=(20,10))
plt.plot(time_scale, waveform_zcr_librosa,color="r")
plt.title("Zero-Cross-rate-librosa")
#plt.show()
bias=waveform_zcr_librosa - waveform_ZCR
#print(waveform_zcr_librosa)

#print(waveform_ZCR)
#print(waveform_STE)

#print(f"the bias is{bias} congratulation!\n")






























